import numpy as np
import matplotlib.pyplot as plt
# import seaborn as sns; sns.set()

from sklearn.decomposition import PCA


def process(src_img, k):
    """
    PCA 算法步骤：

    1. 设有 m行 x n列的矩阵 C 构成的原始数据，记为 n个特征(feature)维度、m个样本数量
    2. 对矩阵 X 作规范化(normalization)处理。这一步相当于令均值 u = 0, 方差 sigma = 1，使得各个元素的概率服从标准正态分布
       2.1 去中心化(zero-center): 对矩阵 X 的每个列元素减去所在列的均值
       2.2 归一化(normalization): 让上一步处理过后的元素值除以所在列的标准差
    3. 计算协方差矩阵 C = X^T @ X / m。矩阵 C 为 n x n 的方阵。
       注：实际上我们是在计算出协方差矩阵 C 的特征值和特征向量后，需要将特征值从大到小排列，并取出前 k 个最大特征值对应的特征向量，
          所以仅需要获得特征值之间的大小关系，对于特征值本身不作要求。而 1/m 会让所有的特征值等比例缩小，但并不会影响特征值的大小关系。
          所以我们在求协方差矩阵 C 时，可以忽略 1/m，即令 C = X^T @ X 即可，方便计算。
    4. 计算协方差矩阵 C 的特征值和特征向量。
    5. 特征向量按照对应特征值从大到小的关系排列成新的特征向量矩阵，取前 k 列组成的矩阵 P(n x k)。
    6. 令 Y = X @ P，则矩阵 Y 即为 n 维矩阵 X 降维到 k 维后的数据，矩阵 P 是这一降维过程的变换矩阵。
       Y 为 m行、k列矩阵，其中包含的数据即为原始数据的“主成分”。

    计算到这一步，在函数体内 return Y 即完成了 PCA 算法。
    而为了让 PCA 算法能够更加直观的展示出来，我们继续利用 PCA 算法对图像进行压缩和重构

    7. 原图像矩阵 X 乘以变换矩阵 P，得到的新矩阵 Y 为 m x k 矩阵，即为原始图像的主成分，相当于对原始对象进行了压缩。
    8. 令 Y 乘以变换矩阵 P 的转置，即 Z = Y @ P^T = X @ P @ P^T = X，这一步骤就完成了图像重构。
       矩阵 P 是由特征向量组成的，为了令协方差等于0,这些特征向量都作了正交单位化处理，即 p1 * p1 = p2 * p2 = ... = 1,
       p1 * p2 = p1 * p3 = p5 * p8 = 0。所以 P @ P^T 得出的是一个单位矩阵 E。

    """
    print("================= Principal Component Analysis ================")

    # 2. 对矩阵 X 作规范化(normalization)处理
    normalized_img, col_means, col_stds = normalization(src_img)

    k_10_vectors = get_k_eigen_vectors(normalized_img, 10)
    pc_by_k_10 = get_principal_component(normalized_img, k_10_vectors)
    img_reconstruct_by_k_10 = img_reconstruct(pc_by_k_10, k_10_vectors, col_means, col_stds)

    k_50_vectors = get_k_eigen_vectors(normalized_img, 50)
    pc_by_k_50 = get_principal_component(normalized_img, k_50_vectors)
    img_reconstruct_by_k_50 = img_reconstruct(pc_by_k_50, k_50_vectors, col_means, col_stds)

    k_200_vectors = get_k_eigen_vectors(normalized_img, 200)
    pc_by_k_200 = get_principal_component(normalized_img, k_200_vectors)
    img_reconstruct_by_k_200 = img_reconstruct(pc_by_k_200, k_200_vectors, col_means, col_stds)

    def __plt_show_all():
        index = 1
        plt.figure(figsize=(20, 20))

        plt.subplot(2, 2, index)
        plt.imshow(src_img, cmap='gray')
        plt.title("original image", color='red')
        index += 1

        plt.subplot(2, 2, index)
        plt.imshow(img_reconstruct_by_k_10, cmap='gray')
        plt.title("k = 10", color='red')
        index += 1

        plt.subplot(2, 2, index)
        plt.imshow(img_reconstruct_by_k_50, cmap='gray')
        plt.title("k = 50", color='red')
        index += 1

        plt.subplot(2, 2, index)
        plt.imshow(img_reconstruct_by_k_200, cmap='gray')
        plt.title("k = 200", color='red')
        index += 1

        plt.show()

    __plt_show_all()


def get_k_eigen_vectors(normalized_img, k=6):
    # 3. 计算协方差矩阵 C = X^T @ X
    COV_C = np.matmul(normalized_img.T, normalized_img)

    # 4. 计算协方差矩阵 C 的特征值和特征向量
    eig_values, eig_vectors = np.linalg.eig(COV_C)

    # 5. 特征向量按照对应特征值从大到小的关系排列成新的特征向量矩阵，取前 k 列组成的矩阵 P(n x k)
    eig_values_descend_index = np.argsort(-eig_values)
    eig_values = eig_values[eig_values_descend_index]
    eig_vectors = eig_vectors[:, eig_values_descend_index]

    return eig_vectors[:, :k]


def get_principal_component(src_matrix, k_vectors):
    return np.matmul(src_matrix, k_vectors)


def normalization(src_img):
    print("============== Normalization ===============")
    # 计算图像每一列的均值
    col_means = np.mean(src_img, axis=0)
    print(col_means)
    # 计算图像每一列的标准差
    col_stds = np.std(src_img, axis=0)
    print(col_stds)

    normalize_img = (src_img - col_means) / col_stds

    return normalize_img, col_means, col_stds


def img_reconstruct(pc_Y, k_vectors, col_means, col_stds):
    pc_Y = np.matmul(pc_Y, k_vectors.T)
    pc_Y = pc_Y * col_stds + col_means
    pc_Y = pc_Y.astype('uint8')

    return pc_Y
